// Copyright Epic Games, Inc. All Rights Reserved.

#include "SmartUIManager.h"
#include "SmartMenuHandle.h"
#include "SmartMenuStyle.h"
#include "SmartCore.h"
#include "ContentBrowserModule.h"
#include "LevelEditor.h"
#include "Framework/MultiBox/MultiBoxBuilder.h"
#include "Framework/MultiBox/MultiBoxExtender.h"
#include "Modules/ModuleManager.h"
#include "ISequencerModule.h"
#include "IPythonScriptPlugin.h"
#include "ToolMenus.h"

#define LOCTEXT_NAMESPACE "SmartMenu"
#define LEVELEDITOR_MODULE_NAME TEXT("LevelEditor")
#define CONTENTBROWSER_MODULE_NAME TEXT("ContentBrowser")
#define SEQUENCER_MODULE_NAME TEXT("Sequencer")


TUniquePtr<FSmartUIManagerImpl> FSmartUIManager::Instance;

class FSmartUIManagerImpl
{
public:
	void Initialize();
	void Shutdown();

private:
	FDelegateHandle ContentBrowserAssetExtenderDelegateHandle;
	FDelegateHandle LevelEditorExtenderDelegateHandle;
	
	void SetupSmartMainMenus(FToolMenuSection& InSection);
	void RegisterMainMenu();
	void SetupSmartToolbarMenus();
	void SetupSmartContextMenus();
	void SetupSmartSequenceMenus();
	void RemoveSmartContextMenus();
	void ActionHook(FString Command, USmartAction* Action);
	void MenuCreator(FMenuBuilder& MenuBuilder, TArray<USmartMenuItemBase*> MenuItems);
	void FillLevelToolbar(FToolBarBuilder& ToolbarBuilder);
	void FillSequenceToolbar(FToolBarBuilder& ToolbarBuilder);
	TSharedRef<SWidget> GenerateSmartToolbarMenu(FString ToolbarType);
	void GenerateSmartMenuContent(FMenuBuilder& MenuBuilder, const TArray<FAssetData>* SelectedAssets, const TArray< AActor*>* SelectedActors, FString MenuPos);
	void GenerateSmartAssetContextMenu(FMenuBuilder& MenuBuilder, TArray<FAssetData> SelectedAssets, FString InMenuType);
	void GenerateSmartActorContextMenu(FMenuBuilder& MenuBuilder, TArray< AActor*> SelectedActors,  FString InMenuType);
	 //Menu extender callbacks
	TSharedRef<FExtender> OnExtendLevelEditor(const TSharedRef<FUICommandList> CommandList, const TArray<AActor*> SelectedActors);
	TSharedRef<FExtender> OnExtendContentBrowserAssetSelectionMenu(const TArray<FAssetData>& SelectedAssets);
};

// SequenceToolbar - 设置工具栏
void FSmartUIManagerImpl::SetupSmartSequenceMenus() {
	const ISequencerModule& SequencerModule = FModuleManager::LoadModuleChecked<ISequencerModule>(SEQUENCER_MODULE_NAME);
	const TSharedPtr<FExtender>SeqToolbarExtender = MakeShareable(new FExtender);
	SeqToolbarExtender->AddToolBarExtension("Base Commands", EExtensionHook::After, nullptr, FToolBarExtensionDelegate::CreateRaw(this, &FSmartUIManagerImpl::FillSequenceToolbar));
	SequencerModule.GetToolBarExtensibilityManager()->AddExtender(SeqToolbarExtender);
}

// SequenceToolbar - 生成工具栏菜单按钮
void FSmartUIManagerImpl::FillSequenceToolbar(FToolBarBuilder & ToolbarBuilder)
{
	ToolbarBuilder.BeginSection(TEXT("Smart"));
	{
		ToolbarBuilder.AddComboButton(
			FUIAction(),
			FOnGetContent::CreateRaw(this, &FSmartUIManagerImpl::GenerateSmartToolbarMenu, FString("Sequencer")),
			LOCTEXT("SmartCombo_Label", "Smart"),
			LOCTEXT("SmartCombo_Tooltip", "Available Smart commands"),
			FSlateIcon(FSmartMenuStyle::GetStyleSetName(), "Smart.Logo")
		);
	}
	ToolbarBuilder.EndSection();
}

TSharedRef<SWidget> FSmartUIManagerImpl::GenerateSmartToolbarMenu(FString ToolbarType) {
	FMenuBuilder MenuBuilder(true, nullptr);
	GenerateSmartMenuContent(MenuBuilder, nullptr, nullptr, ToolbarType);
	return MenuBuilder.MakeWidget();
}

// MainMenu - 生成主菜單
void FSmartUIManagerImpl::SetupSmartMainMenus(FToolMenuSection& InSection)
{
	if (const USmartMenuHandle* MenuHandle = USmartMenuHandle::Get())
	{
		const USmartMenu* MainMenu = MenuHandle->GetMainMenu();
		TArray<USmartMenuItemBase*> MenuItems = MainMenu->Items;
		if (MenuItems.Num() == 0) {
			return;
		}
		for (USmartMenuItemBase* MenuItem : MenuItems)
		{
			if (MenuItem == nullptr) {
				return;
			}
			if (MenuItem->GetType() == EMenuItemType::Separator)
			{
				InSection.AddSeparator(FName(FGuid::NewGuid().ToString()));
			}
			else if (MenuItem->GetType() == EMenuItemType::Action)
			{
				USmartAction* Action = Cast<USmartAction>(MenuItem);
				FString CommandName = Action->Name;
				const FString Command = Action->Command;
				FString Description = Action->Description;
				InSection.AddMenuEntry(
					FName(FGuid::NewGuid().ToString()),
					FText::FromString(CommandName),
					FText::FromString(Description),
					FSlateIcon(),
					FExecuteAction::CreateRaw(this, &FSmartUIManagerImpl::ActionHook, Command, Action)
				);
			}
			else if (MenuItem->GetType() == EMenuItemType::Menu) {
				const USmartMenu* Menu = Cast<USmartMenu>(MenuItem);
				FString CommandName = Menu->Name;
				FString Description = Menu->Description;
				InSection.AddSubMenu(
					FName(FGuid::NewGuid().ToString()),
					FText::FromString(CommandName),
					FText::FromString(Description),
					FNewMenuDelegate::CreateRaw(this, &FSmartUIManagerImpl::MenuCreator, Menu->Items),
					false,
					FSlateIcon()
				);
			}
		}
	}
}

// 創建主菜單
void FSmartUIManagerImpl::RegisterMainMenu()
{
	FToolMenuOwnerScoped OwnerScoped(this);
	UToolMenu* ToolMenu = UToolMenus::Get()->ExtendMenu("LevelEditor.MainMenu");
	if(!ToolMenu->FindSection("SMART"))
	{
		FToolMenuSection& Section = ToolMenu->FindOrAddSection("Help");
		Section.AddMenuEntry(
			FName("SMART"),
			FText::FromString("SMART"),
			FText::FromString("SMART"),
			FSlateIcon(),
			FUIAction()
		);
	}
	
	UToolMenu* SmartToolMenu = UToolMenus::Get()->ExtendMenu("LevelEditor.MainMenu.SMART");
	FToolMenuSection& Section = SmartToolMenu->FindOrAddSection("SMART");
	Section.AddDynamicEntry("SMART", FNewToolMenuSectionDelegate::CreateRaw(this, &FSmartUIManagerImpl::SetupSmartMainMenus));
	// Section.AddDynamicEntry("SMART", FNewToolMenuSectionDelegate::CreateLambda([&](FToolMenuSection& InSection)
	// {
	// 	SetupSmartMainMenus(InSection);
	// }
	// ));
}

// LevelToolbar - 设置视口菜单
void FSmartUIManagerImpl::SetupSmartToolbarMenus()
{
	// 设置图标
	FSmartMenuStyle::SetIcon("Logo", "smart_logo_80px");
	FSmartMenuStyle::SetIcon("ContextLogo", "smart_context_logo");

	// 设置位置：将工具放在 “settings”后面
	FLevelEditorModule& LevelEditorModule = FModuleManager::LoadModuleChecked<FLevelEditorModule>(LEVELEDITOR_MODULE_NAME);
	const TSharedPtr<FExtender> ToolbarExtender = MakeShareable(new FExtender);

	// 生成菜单项
	ToolbarExtender->AddToolBarExtension("Settings", EExtensionHook::After, nullptr, FToolBarExtensionDelegate::CreateRaw(this, &FSmartUIManagerImpl::FillLevelToolbar));
	LevelEditorModule.GetToolBarExtensibilityManager()->AddExtender(ToolbarExtender);
}

// LevelToolbar - 生成视口菜单按钮
void FSmartUIManagerImpl::FillLevelToolbar(FToolBarBuilder& ToolbarBuilder)
{
	ToolbarBuilder.BeginSection(TEXT("Smart"));
	{
		ToolbarBuilder.AddComboButton(
			FUIAction(),
			FOnGetContent::CreateRaw(this, &FSmartUIManagerImpl::GenerateSmartToolbarMenu, FString("Level")),
			LOCTEXT("SmartCombo_Label", "Smart"),
			LOCTEXT("SmartCombo_Tooltip", "Available Smart commands"),
			FSlateIcon(FSmartMenuStyle::GetStyleSetName(), "Smart.Logo")
		);
	}
	ToolbarBuilder.EndSection();
}

void FSmartUIManagerImpl::SetupSmartContextMenus()
{
	// Register Content Browser menu extender
	FContentBrowserModule& ContentBrowserModule = FModuleManager::LoadModuleChecked<FContentBrowserModule>(CONTENTBROWSER_MODULE_NAME);

	TArray<FContentBrowserMenuExtender_SelectedAssets>& CBAssetMenuExtenderDelegates = ContentBrowserModule.GetAllAssetViewContextMenuExtenders();
	CBAssetMenuExtenderDelegates.Insert(FContentBrowserMenuExtender_SelectedAssets::CreateRaw(this, &FSmartUIManagerImpl::OnExtendContentBrowserAssetSelectionMenu),0);
	//CBAssetMenuExtenderDelegates.Add(FContentBrowserMenuExtender_SelectedAssets::CreateRaw(this, &FSmartUIManagerImpl::OnExtendContentBrowserAssetSelectionMenu));
	ContentBrowserAssetExtenderDelegateHandle = CBAssetMenuExtenderDelegates.Top().GetHandle();

	// Register Level Editor menu extender
	FLevelEditorModule& LevelEditorModule = FModuleManager::GetModuleChecked<FLevelEditorModule>(LEVELEDITOR_MODULE_NAME);

	TArray<FLevelEditorModule::FLevelViewportMenuExtender_SelectedActors>& LevelEditorMenuExtenderDelegates = LevelEditorModule.GetAllLevelViewportContextMenuExtenders();
	LevelEditorMenuExtenderDelegates.Insert(FLevelEditorModule::FLevelViewportMenuExtender_SelectedActors::CreateRaw(this, &FSmartUIManagerImpl::OnExtendLevelEditor), 0);
	//LevelEditorMenuExtenderDelegates.Add(FLevelEditorModule::FLevelViewportMenuExtender_SelectedActors::CreateRaw(this, &FSmartUIManagerImpl::OnExtendLevelEditor));
	LevelEditorExtenderDelegateHandle = LevelEditorMenuExtenderDelegates.Top().GetHandle();
}

void FSmartUIManagerImpl::RemoveSmartContextMenus()
{
	if (FModuleManager::Get().IsModuleLoaded(LEVELEDITOR_MODULE_NAME))
	{
		FLevelEditorModule& LevelEditorModule = FModuleManager::GetModuleChecked<FLevelEditorModule>(LEVELEDITOR_MODULE_NAME);
		TArray<FLevelEditorModule::FLevelViewportMenuExtender_SelectedActors>& LevelEditorMenuExtenderDelegates = LevelEditorModule.GetAllLevelViewportContextMenuExtenders();
		LevelEditorMenuExtenderDelegates.RemoveAll([this](const FLevelEditorModule::FLevelViewportMenuExtender_SelectedActors& Delegate) { return Delegate.GetHandle() == LevelEditorExtenderDelegateHandle; });
	}
	if (FModuleManager::Get().IsModuleLoaded(CONTENTBROWSER_MODULE_NAME))
	{
		FContentBrowserModule& ContentBrowserModule = FModuleManager::GetModuleChecked<FContentBrowserModule>(CONTENTBROWSER_MODULE_NAME);
		TArray<FContentBrowserMenuExtender_SelectedAssets>& CBAssetMenuExtenderDelegates = ContentBrowserModule.GetAllAssetViewContextMenuExtenders();
		CBAssetMenuExtenderDelegates.RemoveAll([this](const FContentBrowserMenuExtender_SelectedAssets& Delegate) { return Delegate.GetHandle() == ContentBrowserAssetExtenderDelegateHandle; });
	}
}

// 根据类型生成菜单
void FSmartUIManagerImpl::GenerateSmartMenuContent(FMenuBuilder& MenuBuilder, const TArray<FAssetData>* SelectedAssets, const TArray< AActor*>* SelectedActors, FString MenuPos)
{
	if (USmartMenuHandle* MenuHandle = USmartMenuHandle::Get())
	{
		USmartMenu* MainMenu;
		TArray<USmartMenuItemBase*> EnvMenuItems;
		if (MenuPos == "Level") {
			MainMenu = MenuHandle->GetLevelToolbarMenu();
			//EnvMenuItems = MenuHandle->GetMenuItemsFromEnv(FString("UNREAL_SMART_LEVEL_MENU_CONFIGS"));
		}
		else if (MenuPos == "Asset") {
			MainMenu = MenuHandle->GetAssetMenu();
			//EnvMenuItems = MenuHandle->GetMenuItemsFromEnv(FString("UNREAL_SMART_ASSET_MENU_CONFIGS"));
		}
		else if (MenuPos == "AssetSmart") {
			MainMenu = MenuHandle->GetAssetSmartMenu();
			//EnvMenuItems = MenuHandle->GetMenuItemsFromEnv(FString("UNREAL_SMART_ASSET_MENU_CONFIGS"));
		}
		else if (MenuPos == "Actor") {
			MainMenu = MenuHandle->GetActorMenu();
			//EnvMenuItems = MenuHandle->GetMenuItemsFromEnv(FString("UNREAL_SMART_ACTOR_MENU_CONFIGS"));
		}
		else if (MenuPos == "ActorSmart") {
			MainMenu = MenuHandle->GetActorSmartMenu();
			//EnvMenuItems = MenuHandle->GetMenuItemsFromEnv(FString("UNREAL_SMART_ACTOR_MENU_CONFIGS"));
		}
		else if (MenuPos == "Sequencer") {
			MainMenu = MenuHandle->GetSequencerToolbarMenu();
			//EnvMenuItems = MenuHandle->GetMenuItemsFromEnv(FString("UNREAL_SMART_SEQUENCER_MENU_CONFIGS"));
		}
		else {
			return;
		}
		//if (EnvMenuItems.Num() != 0) {
		//	MainMenu->Items.Append(EnvMenuItems);
		//}
		MenuCreator(MenuBuilder, MainMenu->Items);
	}
}

// 菜单生成器
void FSmartUIManagerImpl::MenuCreator(FMenuBuilder& MenuBuilder, TArray<USmartMenuItemBase*> MenuItems){
	if (MenuItems.Num() == 0) {
		return;
	}
	for (USmartMenuItemBase* MenuItem : MenuItems)
	{
		if (MenuItem == nullptr) {
			return;
		}
		if (MenuItem->GetType() == EMenuItemType::BeginSection)
		{
			const USmartBeginSection* BeginSection = Cast<USmartBeginSection>(MenuItem);
			FString CommandName = BeginSection->Name;
			MenuBuilder.BeginSection(NAME_None, FText::FromString(CommandName));
		}
		else if (MenuItem->GetType() == EMenuItemType::EndSection)
		{
			USmartEndSection* EndSection = Cast<USmartEndSection>(MenuItem);
			MenuBuilder.EndSection();
		}
		else if (MenuItem->GetType() == EMenuItemType::Separator)
		{
			MenuBuilder.AddMenuSeparator();
		}
		else if (MenuItem->GetType() == EMenuItemType::Action)
		{
			USmartAction* Action = Cast<USmartAction>(MenuItem);
			FString CommandName = Action->Name;
			const FString Command = Action->Command;
			FString Description = Action->Description;
			
			MenuBuilder.AddMenuEntry(
				FText::FromString(CommandName),
				FText::FromString(Description),
				FSlateIcon(),
				FExecuteAction::CreateRaw(this, &FSmartUIManagerImpl::ActionHook, Command, Action)
			);
		}
		else if (MenuItem->GetType() == EMenuItemType::Menu) {
			const USmartMenu* Menu = Cast<USmartMenu>(MenuItem);
			FString CommandName = Menu->Name;
			FString Description = Menu->Description;
			MenuBuilder.AddSubMenu(
				FText::FromString(CommandName),
				FText::FromString(Description),
				FNewMenuDelegate::CreateRaw(this, &FSmartUIManagerImpl::MenuCreator, Menu->Items),
				false,
				FSlateIcon()
			);
		}
	}
}

void FSmartUIManagerImpl::ActionHook(FString Command, USmartAction* Action)
{
	if(Action->bIsWrapper)
	{
		Action->ExecuteCommand();
	}else
	{
		IPythonScriptPlugin::Get()->ExecPythonCommand(*Command);
	}
}

void FSmartUIManagerImpl::GenerateSmartAssetContextMenu(FMenuBuilder& MenuBuilder, TArray<FAssetData> SelectedAssets, FString InMenuType)
{
	GenerateSmartMenuContent(MenuBuilder, &SelectedAssets, nullptr, InMenuType);
}

void FSmartUIManagerImpl::GenerateSmartActorContextMenu(FMenuBuilder& MenuBuilder, TArray<AActor*> SelectedActors, FString InMenuType)
{
	GenerateSmartMenuContent(MenuBuilder, nullptr, &SelectedActors, InMenuType);
}

TSharedRef<FExtender> FSmartUIManagerImpl::OnExtendContentBrowserAssetSelectionMenu(const TArray<FAssetData>& SelectedAssets)
{
	// Menu extender for Content Browser context menu when an asset is selected
	TSharedRef<FExtender> Extender(new FExtender());
	if (SelectedAssets.Num() > 0)
	{
		Extender->AddMenuExtension("AssetContextReferences", EExtensionHook::After, nullptr, FMenuExtensionDelegate::CreateLambda(
			[this, SelectedAssets](FMenuBuilder& MenuBuilder)
			{
				MenuBuilder.BeginSection(NAME_None, FText::FromString("Smart Tools"));
				// 二级菜单（子菜单）
				MenuBuilder.AddSubMenu(
					LOCTEXT("Smart_ContextMenu", "Smart"),
					FText(),
					FNewMenuDelegate::CreateRaw(this, &FSmartUIManagerImpl::GenerateSmartAssetContextMenu, SelectedAssets, FString("AssetSmart")),
					false,
					FSlateIcon(FSmartMenuStyle::GetStyleSetName(), "Smart.ContextLogo")
				);
				// 一级菜单项
				GenerateSmartAssetContextMenu(MenuBuilder, SelectedAssets, FString("Asset"));
				MenuBuilder.EndSection();
			}));
	}
	return Extender;
}

TSharedRef<FExtender> FSmartUIManagerImpl::OnExtendLevelEditor(const TSharedRef<FUICommandList> CommandList, const TArray<AActor*> SelectedActors)
{
	// Menu extender for Level Editor and World Outliner context menus when an actor is selected
	TSharedRef<FExtender> Extender(new FExtender());
	if (SelectedActors.Num() > 0)
	{
		Extender->AddMenuExtension("LevelViewportAttach", EExtensionHook::After, nullptr, FMenuExtensionDelegate::CreateLambda(
			[this, SelectedActors](FMenuBuilder& MenuBuilder)
			{
				MenuBuilder.BeginSection(NAME_None, FText::FromString("Smart Tools"));
				MenuBuilder.AddSubMenu(
					LOCTEXT("Smart_ContextMenu", "Smart"),
					FText(),
					FNewMenuDelegate::CreateRaw(this, &FSmartUIManagerImpl::GenerateSmartActorContextMenu, SelectedActors, FString("ActorSmart")),
					false,
					FSlateIcon(FSmartMenuStyle::GetStyleSetName(), "Smart.ContextLogo")
				);
				GenerateSmartActorContextMenu(MenuBuilder, SelectedActors, FString("Actor"));
				MenuBuilder.EndSection();
			}));
	}
	return Extender;
}

// 初始化
void FSmartUIManagerImpl::Initialize()
{
	FSmartMenuStyle::Initialize();
	RegisterMainMenu();
	SetupSmartToolbarMenus();
	SetupSmartContextMenus();
	SetupSmartSequenceMenus();
	//FCoreDelegates::OnPostEngineInit.AddRaw(this, &FSmartUIManagerImpl::userScript);
	//FCoreDelegates::OnFEngineLoopInitComplete.AddRaw(this, USmartCore::userScript);
	//FCoreDelegates::OnFEngineLoopInitComplete.AddLambda(USmartCore::userScript);
	FCoreDelegates::OnFEngineLoopInitComplete.AddStatic(USmartCore::UserScript);
	// UToolMenus::RegisterStartupCallback(FSimpleMulticastDelegate::FDelegate::CreateRaw(this, &FSmartUIManagerImpl::RegisterMenus));
}

void FSmartUIManagerImpl::Shutdown()
{
	RemoveSmartContextMenus();
	FSmartMenuStyle::Shutdown();
}

void FSmartUIManager::Initialize()
{
	if (!Instance.IsValid())
	{
		Instance = MakeUnique<FSmartUIManagerImpl>();
		Instance->Initialize();
	}

}

void FSmartUIManager::Shutdown()
{
	if (Instance.IsValid())
	{
		Instance->Shutdown();
		Instance.Reset();
	}
}

#undef LOCTEXT_NAMESPACE
